package org.softee.concurrent;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

public class ConcurrentProcessor<I, O> {
    private static final int AVAILABLE_PROCESSORS = Runtime.getRuntime().availableProcessors();
    private final ExecutorService service;
    private final Function<I, O> taskFunction;
    private final boolean removeNullResults;

    public ConcurrentProcessor(Function<I, O> taskFunction) {
        this(taskFunction, Executors.newFixedThreadPool(AVAILABLE_PROCESSORS * 2));
    }

    public ConcurrentProcessor(Function<I, O> taskFunction, ExecutorService service) {
        this(taskFunction, service, false);
    }

    public ConcurrentProcessor(Function<I, O> taskFunction, ExecutorService executor, boolean removeNullResults) {
        this.service = executor;
        this.taskFunction = taskFunction;
        this.removeNullResults = removeNullResults;
    }


    public List<O> process(Collection<I> sources) throws InterruptedException, ExecutionException {
        int count = sources.size();
        final Queue<Future<O>> futures = new ArrayDeque<Future<O>>(count);

        for (final I source : sources) {
            final Callable<O> callable = new Callable<O>() {
                public O call() {
                    return taskFunction.apply(source);
                }
            };

            futures.add(service.submit(callable));
        }

        try {
            List<O> results = new ArrayList<O>(count);
            for (int i = 0; i < count; i++) {
                Future<O> future = futures.remove();
                O result = future.get();
                if (result != null || !removeNullResults) {
                    results.add(result);
                }
            }
            return results;
        } finally {
            cancel(futures);
        }
    }

    private void cancel(Collection<Future<O>> futures) {
        for (Future<O> future : futures) {
            future.cancel(false);
        }
    }

    public static interface Function<I, O> {
        O apply(I input);
    }
}
